from ultralytics import YOLO

# 关键：Windows 多进程必须加这行保护，避免重复创建进程
if __name__ == '__main__':
    # 加载模型
    model = YOLO("yolo11n.pt")  # 预训练模型

    # 开始训练（参数可根据你的需求调整）
    results = model.train(
        epochs=100,  # 训练轮数（先从10轮测试，避免之前的100轮耗时过长）
        data="C:/Users/Lenovo/Desktop/ultralytics-main/ultralytics/cfg/datasets/coco8.yaml",  # 你的数据集配置文件
        batch=8,  # 批次大小：RTX 3050 Laptop GPU 显存4G，建议从8开始（之前16可能略满）
        imgsz=640,  # 输入图像尺寸（YOLO默认，无需修改）
        device=0,  # 使用第1块GPU（0表示GPU，-1表示CPU）
        workers=4,  # 数据加载线程数：Windows建议设为4（之前8可能过多，导致进程冲突）
        project="herb_detection",  # 结果保存文件夹
        name="exp_test",  # 本次实验名称（避免覆盖之前的结果）
        pretrained=True  # 使用预训练权重（加速收敛）
    )